package com.bruce.ai.alibaba.workflow.controller;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.http.MediaType;
import org.springframework.http.codec.ServerSentEvent;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;

import com.alibaba.cloud.ai.graph.CompileConfig;
import com.alibaba.cloud.ai.graph.CompiledGraph;
import com.alibaba.cloud.ai.graph.OverAllState;
import com.alibaba.cloud.ai.graph.StateGraph;
import com.alibaba.cloud.ai.graph.checkpoint.config.SaverConfig;
import com.alibaba.cloud.ai.graph.checkpoint.savers.MemorySaver;

import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

/**
 * 短期记忆工作流
 */
@RestController
@RequestMapping("/short-memory")
public class ShortMemoryWorkflowController {

    private CompiledGraph memoryWorkflowCompiledGraph ;



    private static final Logger logger = LoggerFactory.getLogger(ShortMemoryWorkflowController.class);


    public ShortMemoryWorkflowController(@Qualifier("memoryWorkflow") StateGraph memoryWorkflow) throws Exception {
        // 定义检查点配置 内存检查点
        SaverConfig saverConfig = SaverConfig.builder().register(new MemorySaver())
                .build();
        this.memoryWorkflowCompiledGraph = memoryWorkflow
                .compile(CompileConfig.builder().saverConfig(saverConfig).build());
    }

    /*
     * 展开工作流
     */
    @RequestMapping(value ="/expand", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<ServerSentEvent<String>> expand(@RequestParam("messages") String messages) throws Exception{
        List<Map<String, String>> messageMaps = new ArrayList<>();
        messageMaps.add(Map.of("role", "user", "content", messages));
        return Mono.defer(() -> Mono.justOrEmpty(memoryWorkflowCompiledGraph.invoke(Map.of("messages", messageMaps)))
                .map(OverAllState::data)
                .map(data -> ServerSentEvent.builder(data.toString()).build()))
            .flux()
            .doOnCancel(() -> logger.info("Client disconnected from stream"))
            .doOnError(e -> logger.error("Error occurred during streaming", e));
    }

}
